#include <iostream>
#include <WINSOCK2.h>
#include <time.h>
#include <string.h>
#include <fstream>
#pragma comment(lib, "ws2_32.lib")

using namespace std;

SOCKADDR_IN server_addr;
SOCKET server;
int len; //sizeof(server_addr)
const int port = 1234;
const int timeout = CLOCKS_PER_SEC;
const int bufsize = 1024;//缓冲区长度

const unsigned char SYN = 0x1; //SYN = 1 ACK = 0
const unsigned char ACK = 0x2;//SYN = 0, ACK = 1
const unsigned char SYN_ACK = 0x3;//SYN = 1, ACK = 1
const unsigned char FIN = 0x4;//FIN = 1, ACK = 0
const unsigned char FIN_ACK = 0x5;//FIN = 1 ACK = 0
const unsigned char OVER = 0x7;//结束标志

struct HEADER
{
    uint16_t sum = 0; // 校验和
    uint16_t len = 0; //数据长度
    unsigned char flag = 0; //标记位_ _ _ _ _ FIN ACK SYN 
    unsigned char seq = 0; //序列号，0~255，超出取模
};

void open_Socket()
{
    WSADATA wsadata;
    WSAStartup(MAKEWORD(2, 2), &wsadata);

    server_addr.sin_family = AF_INET;//使用IPV4
    server_addr.sin_port = htons(port);
    server_addr.sin_addr.s_addr = htonl(2130706433);

    server = socket(AF_INET, SOCK_DGRAM, 0);
    bind(server, (SOCKADDR*)&server_addr, sizeof(server_addr));
    cout << "等待客户端连接\n";
}

uint16_t cksum(uint16_t* buffer, int count) {
    /*if (count == 0)
        return ~(0);*/
    uint16_t* buf = (uint16_t*)malloc(count);
    memcpy(buf, buffer, count);
    count = (count + 1) / 2;
    register uint32_t sum = 0;
    while (count--) {
        sum += *buf++;
        if (sum & 0XFFFF0000) {
            sum &= 0XFFFF;
            sum++;
        }
    }
    return ~(sum & 0XFFFF);
}


bool shakehand()
{

    HEADER head;
    char* buf = new char[sizeof(head)];

    //第一次握手
    while (1) {
        if (recvfrom(server, buf, sizeof(head), 0, (sockaddr*)&server_addr, &len) == -1)
            return false;
        memcpy(&head, buf, sizeof(head));
        if (head.flag == SYN && cksum((uint16_t*)&head, sizeof(head)) == 0)
        {
            cout << "第一次握手成功\n";
            break;
        }
    }

    //发送第二次握手信息
    head.sum = 0;
    head.flag = ACK;
    head.sum = cksum((uint16_t*)&head, sizeof(head));
    memcpy(buf, &head, sizeof(head));
    if (sendto(server, buf, sizeof(head), 0, (sockaddr*)&server_addr, len) == -1)
        return false;
    clock_t start = clock();//记录第二次握手发送时间

    //接收第三次握手
    while (recvfrom(server, buf, sizeof(head), 0, (sockaddr*)&server_addr, &len) <= 0)
    {
        if (clock() - start > timeout)//超时重传
        {
            head.flag = ACK;
            head.sum = 0;
            head.flag = cksum((uint16_t*)&head, sizeof(head));
            memcpy(buf, &head, sizeof(head));
            if (sendto(server, buf, sizeof(head), 0, (sockaddr*)&server_addr, len) == -1)
                return -1;
            cout << "第二次握手超时，重传ing\n";
        }
    }

    HEADER temp1;
    memcpy(&temp1, buf, sizeof(head));
    if (temp1.flag == SYN_ACK && cksum((uint16_t*)&temp1, sizeof(temp1) == 0))
        cout << "三次握手完成,连接成功\n";
    else
    {
        cout << "连接失败\n";
        return false;
    }
    return true;
}

int _recv(char* message)
{
    HEADER head;
    char* buf = new char[bufsize + sizeof(head)];
    uint32_t mlen = 0;//文件长度
    int seq = 0;
    int index = 0;

    while (1)
    {
        int length = recvfrom(server, buf, sizeof(head) + bufsize, 0, (sockaddr*)&server_addr, &len); //报文长度
        memcpy(&head, buf, sizeof(head));
        if (head.flag == OVER && cksum((uint16_t*)&head, sizeof(head)) == 0){//结束
            cout << "文件接收完成" << endl;
            break;
        }
        if (head.flag == unsigned char(0) && cksum((uint16_t*)buf, length - sizeof(head))){
            if (seq != int(head.seq)){ //发错了，返回ACK
                head.flag = ACK;
                head.len = 0;
                head.seq = (unsigned char)seq;
                head.sum = 0;
                head.sum = cksum((uint16_t*)&head, sizeof(head));
                memcpy(buf, &head, sizeof(head));
                //重传ACK
                sendto(server, buf, sizeof(head), 0, (sockaddr*)&server_addr, len);
                cout << "发送 ACK:" << (int)head.seq << " SEQ:" << (int)head.seq << endl;
                continue;//丢弃该数据包
            }
            seq = int(head.seq);
            if (seq > 255)
                seq = seq - 256;
            cout << "发送数据: " << length - sizeof(head) << " bytes, FlAG:" << int(head.flag) << " SEQ : " << int(head.seq) << " SUM:" << int(head.sum) << endl;
            char* temp = new char[length - sizeof(head)];
            memcpy(temp, buf + sizeof(head), length - sizeof(head));
            memcpy(message + mlen, temp, length - sizeof(head));
            mlen = mlen + int(head.len);

            //返回ACK
            head.flag = ACK;
            head.len = 0;
            head.seq = (unsigned char)seq;
            head.sum = 0;
            uint16_t temp1 = cksum((uint16_t*)&head, sizeof(head));
            head.sum = temp1;
            memcpy(buf, &head, sizeof(head));
            //重发该包的ACK
            sendto(server, buf, sizeof(head), 0, (sockaddr*)&server_addr, len);
            cout << "发送 ACK:" << (int)head.seq << " SEQ:" << (int)head.seq << endl;
            seq++;
            if (seq > 255)
            {
                seq = seq - 256;
            }
        }
    }
    //发送OVER信息
    head.sum = 0;
    head.flag = OVER;
    head.sum = cksum((uint16_t*)&head, sizeof(head));
    memcpy(buf, &head, sizeof(head));
    sendto(server, buf, sizeof(head), 0, (sockaddr*)&server_addr, len);
    message[mlen] = 0;
    return mlen;
}

int wavehand()
{
    HEADER head;
    char* buf = new char[sizeof(head)];
    while (1)
    {
        int length = recvfrom(server, buf, sizeof(head) + bufsize, 0, (sockaddr*)&server_addr, &len);//接收报文长度
        memcpy(&head, buf, sizeof(head));
        if (head.flag == FIN && cksum((uint16_t*)&head, sizeof(head)) == 0)
        {
            cout << "第一次挥手成功" << endl;
            break;
        }
    }
    //发送第二次挥手信息
    head.flag = ACK;
    head.sum = 0;
    uint16_t temp = cksum((uint16_t*)&head, sizeof(head));
    head.sum = temp;
    memcpy(buf, &head, sizeof(head));
    if (sendto(server, buf, sizeof(head), 0, (sockaddr*)&server_addr, len) == -1)
        return -1;
    clock_t start = clock();

    //接收第三次挥手
    while (recvfrom(server, buf, sizeof(head), 0, (sockaddr*)&server_addr, &len) <= 0)
    {
        if (clock() - start > timeout)
        {
            head.flag = ACK;
            head.sum = 0;
            uint16_t temp = cksum((uint16_t*)&head, sizeof(head));
            head.flag = temp;
            memcpy(buf, &head, sizeof(head));
            if (sendto(server, buf, sizeof(head), 0, (sockaddr*)&server_addr, len) == -1)
                return -1;
            cout << "第二次挥手超时重传ing\n";
        }
    }

    HEADER temp1;
    memcpy(&temp1, buf, sizeof(head));
    if (temp1.flag == FIN_ACK && cksum((uint16_t*)&temp1, sizeof(temp1) == 0))
        cout << "第三次挥手成功\n";
    else
    {
        cout << "第三次挥手error\n";
        return -1;
    }

    //发送第四次挥手信息
    head.flag = FIN_ACK;
    head.sum = 0;
    temp = cksum((uint16_t*)&head, sizeof(head));
    head.sum = temp;
    memcpy(buf, &head, sizeof(head));
    if (sendto(server, buf, sizeof(head), 0, (sockaddr*)&server_addr, len) == -1)
    {
        cout << "第四次挥手error\n";
        return -1;
    }
    cout << "四次挥手结束，连接断开！\n";
    return 1;
}


int main()
{
    open_Socket();
    len = sizeof(server_addr);
    //建立连接
    if (!shakehand())
        return 0;
    char* name = new char[20];
    char* data = new char[100000000];
    int n = _recv(name);
    int d = _recv(data);
    string output;
    for (int i = 0; i < n; i++)
        output += name[i];
    wavehand();
    ofstream ofile(output.c_str(), ofstream::binary);
    for (int i = 0; i < d; i++)
        ofile << data[i];

    cout << "文件接收成功\n";
    ofile.close();
    system("pause");
    return 0;
}
